import os
import decode
import encode
import psnr
import time
import ssim

save_yuv_result=0
save_mp4_result=0

if __name__ == "__main__":
	print("---compute_quiality.py--")
	#1st decode
	video_path= "/home/lzh/Videos/600_frames.h264"
	total_frames=300
	original_file="original.yuv"
	width=1920
	height=1080
	exit_code = decode.decode_cmd(video_path,original_file,total_frames,width,height)
	assert(exit_code == 0)
    #2nd encode 
	encode_parameters=[{"bitrate":2000000,"preset":"superfast","enc_params":"low_delay=1"},
                   {"bitrate":1000000,"preset":"superfast","enc_params":"low_delay=1"},
                   {"bitrate":500000,"preset":"superfast","enc_params":"low_delay=1"}]
	#encode_parameters=[{"bitrate":2000000,"preset":"superfast","enc_params":"low_delay=1"}]


	for parameter in encode_parameters:
		encode_out_file = str(parameter["bitrate"])+".mp4"
		exit_code = encode.encode_cmd(original_file,width,height,parameter["bitrate"],parameter["preset"],parameter["enc_params"],encode_out_file)
		
		assert(exit_code==0)
		#3rd decode	
		decode_out_file=str(parameter["bitrate"])+".yuv"
		exit_code=decode.decode_cmd(encode_out_file,decode_out_file,total_frames,width,height)
		assert(exit_code ==0)

    	#4th psnr
		psnr_values, avg_psnr = psnr.process_yuv_files(width, height,decode_out_file,original_file, total_frames)
		#5th ssim
		ssim.calculate_yuv_video_ssim(decode_out_file, original_file, width, height, total_frames)
		#at the end
		if(save_yuv_result ==0):
			os.system("rm -rf "+ decode_out_file) 
		if(save_mp4_result == 0):
			os.system("rm -rf "+ encode_out_file)
		
	if(save_yuv_result ==0):
		os.system("rm -rf "+ original_file)
